package nachos.vm;

import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Random;

import nachos.machine.*;
import nachos.threads.*;
import nachos.userprog.*;

/**
 * A <tt>UserProcess</tt> that supports demand-paging.
 */
public class VMProcess extends UserProcess {
    /**
     * Allocate a new process.
     */
    public VMProcess() {
    super();
    totalStackPages = 0;
    }

    /**
     * Save the state of this process in preparation for a context switch.
     * Called by <tt>UThread.saveState()</tt>.
     */
    public void saveState() {
	super.saveState();
    }

    /**
     * Restore the state of this process after a context switch. Called by
     * <tt>UThread.restoreState()</tt>.
     */
    public void restoreState() {
    	Lib.debug(dbgVM, "Restoring state after context switch");
        //group 7
        //invalidate all the TLB entries on context switch
        //do this by setting valid flag to false.
        for(int i = 0; i < Machine.processor().getTLBSize();i++){
              TranslationEntry te = Machine.processor().readTLBEntry(i);
              te.valid = false;
              Machine.processor().writeTLBEntry(i,te);
        }
    }

    /**
     * Initializes page tables for this process so that the executable can be
     * demand-paged.
     *
     * @return	<tt>true</tt> if successful.
     */
    protected boolean loadSections() {
    	Lib.debug(dbgProcess, "in vm: sections: " + coff.getNumSections());
    	
    	for (int i=0;i< coff.getNumSections(); i++) {
		    CoffSection section = coff.getSection(i);
		    
		    int firstVpn = section.getFirstVPN();
		    Lib.debug(dbgVM, "first vpn: " + firstVpn + ", section length: " + section.getLength());		
		    //handle each page. The first page address is 
		    // the vpn. The last page address is vpn+length -1. 
		    // so set j to firstVPN.
		    // the page number will be the index in the table, 
		    // the section number will be the value. 
		    for(int j=firstVpn;j<firstVpn + section.getLength();j++){
		 	   pageSectionMap.put(j, i);
		    }
		}

		return true;
    }
     
    
    /* Group 7
     * 
	 * Handles the exec() system call.
	 * 
	 * Checks that the number of arguments (numArgs) is greater than or equal to 0.
	 * Checks that this is a valid program name. 
	 * Runs process, returns the processId of the child
	 * 
	 */
    
    protected int handleExec(int f, int numArgs, int a) {
    	//assert that numArgs is non-negative
    	String progName = this.readVirtualMemoryString(f, ADDRESS_LENGTH);
    	if(progName == null) return -1;
    	
    	//args
    	if(numArgs < 0) return -1;
    	
    	String[] args = new String[numArgs];
    	int addr = a;
    	for(int i=0;i<args.length;i++){
    	    byte[] data = new byte[4]; 
    	    int bytesRead = this.readVirtualMemory(addr, data);
    	    int ptrArgv = Lib.bytesToInt(data, 0);
    	    args[i] = this.readVirtualMemoryString(ptrArgv, ADDRESS_LENGTH);
    	    System.out.println("args: " + args[i]);
    	    addr += 4;
    	}
    	
    	VMProcess child = new VMProcess();
    	children.add(child);
    	child.parent = this;

    	
    	return child.processNumber;
    }

     
  
     
    public int handleSyscall(int syscall, int a0, int a1, int a2, int a3) {
    	//Lib.debug(dbgProcess, "syscall " + syscall);
    	  switch (syscall) {
    		case syscallExec:
    			return handleExec(a0,a1,a2);
    		case syscallExit:
    		case syscallHalt:
    			System.out.println("Process # " +  processNumber + 
    					": page faults " + pageFaults);
    		default:
    			return super.handleSyscall(syscall, a0, a1, a2, a3);
    	  }
    }
    
    /**
     * Release any resources allocated by <tt>loadSections()</tt>.
     */
    protected void unloadSections() {
    	Lib.debug(dbgVM, "Unloading the sections");
        for(int i = 0; i < numPages; i++){
           VMKernel.removeReferencesToPage(this.processNumber, i);
        }
        Lib.debug(dbgVM, "Unloading done");
    }    

    /**
     * Handle a user exception. Called by
     * <tt>UserKernel.exceptionHandler()</tt>. The
     * <i>cause</i> argument identifies which exception occurred; see the
     * <tt>Processor.exceptionZZZ</tt> constants.
     *
     * @param	cause	the user exception that occurred.
     */
    public void handleException(int cause) {
	Processor processor = Machine.processor();

	switch (cause) {
           //Group 7 - Handle TLB Miss
            case Processor.exceptionTLBMiss:
                handleTLBMiss(processor);
                break;


	default:
	    super.handleException(cause);
	    break;
	}
    }
    public int writeVirtualMemory(int vaddr, byte[] data) {
	return writeVirtualMemory(vaddr, data, 0, data.length);
    }
    //Group 7
    //Overriding writeVirtualMemory to set dirty and used bits
     public int writeVirtualMemory(int vaddr, byte[] data, int offset,
				  int length) {
    	 Lib.debug(dbgVM, "----- In writeVirtualMemory writing this many bytes:  " + length );
    	 Lib.assertTrue(offset >= 0 && length >= 0 && offset+length <= data.length);
    	    	
 		int amountWritten = 0;
 		byte[] memory = Machine.processor().getMemory();
 		int vpn = 0;
 		
 		try{ 			
 			while(amountWritten < length){
 				vpn = Processor.pageFromAddress(vaddr);
 				Lib.debug(dbgVM, "Virtual address to write to: " + vaddr + ", the virtual page to write to: " + vpn);
 				 //first make sure this page is in memory:
 		    	 if (!VMKernel.isInMemory(processNumber, vpn)){
 		      		handlePageFault(vpn);
 		      	} 
 		    	 //now it should be in memory, loaded by the handlePageFault.
 		    	 
 				//mark the page as in use:
 	 		    VMKernel.pinPage(this.processNumber, vpn);
 	 		     	 				
 				//get the physical address:
 				int paddr = getPhysicalAddress(vaddr);
 				if (paddr < 0 || paddr >= memory.length){
 					Lib.debug(dbgProcess, "returning early from write, physical address was: "
 								+ paddr + ", memory length is " + memory.length);
 					VMKernel.unpinPage(this.processNumber, vpn);
 				    return amountWritten;
 				}
 				//we should only write until the end of the page, as we do 
 				// not know where the next physical page may be. Will have
 				//to recalculate the virtual address after that. 
 				int amount = Math.min(length, pageSize - length);
 				System.arraycopy(data, offset, memory, paddr, amount);
 				amountWritten = amountWritten + amount;
 				vaddr = vaddr + amount;
 				length = length - amount;
 				offset = offset + amount;
 				
 				for(int i = 0; i < Machine.processor().getTLBSize(); i++){
 		              TranslationEntry te = Machine.processor().readTLBEntry(i);
 		              if(te.vpn == vpn){
 		                  te.dirty = true;
 		                  te.used = true;
 		                  Machine.processor().writeTLBEntry(i, te);
 		              }
 		         }
 				String key = VMKernel.createKey(this.processNumber, vpn);
 				VMKernel.getPageFromMemory(key).dirty = true;
 				VMKernel.getPageFromMemory(key).used = true;
 				Lib.debug(dbgVM, "Set dirty and used bits on page at: " + key +
 							" at ppn: " + VMKernel.getPageFromMemory(key).ppn);
 				
 				
 	    	}
 		}catch(IOException e){
 			Lib.debug(dbgProcess, "IOException in getPhysicalAddress: " + e);
 		}
 		 
 		VMKernel.unpinPage(this.processNumber, vpn);
 		Lib.debug(dbgVM, "----- exiting writeVirtualMemory");
        return amountWritten;
     }

     public int readVirtualMemory(int vaddr, byte[] data) {
	return readVirtualMemory(vaddr, data, 0, data.length);
    }

    //Group 7
    //Overriding readVirtualMemory to set  used bit
     public int readVirtualMemory(int vaddr, byte[] data, int offset,
				 int length){
    	Lib.debug(dbgVM, "----- In readVirtualMemory");
    	int amountRead = 0;
 		byte[] memory = Machine.processor().getMemory();
 		int vpn=0;
 		
 		try{
 			while(amountRead < length){
 				//get the physical address:
 				vpn = Processor.pageFromAddress(vaddr);
 				 //first make sure this page is in memory:
 			    if (!VMKernel.isInMemory(processNumber, vpn)){
 			      	handlePageFault(vpn);
 			     } 
 			    //now it should be in memory, loaded by the handlePageFault.
 			    //mark the page as in use:
 	 		    VMKernel.pinPage(this.processNumber, vpn);
 	 		   	 				
 				//get the physical address:
 				int paddr = getPhysicalAddress(vaddr);
 				if (paddr < 0 || paddr >= memory.length){
 					Lib.debug(dbgProcess, "returning early from read, physical address was: "
 								+ paddr + ", memory length is " + memory.length);
 					VMKernel.unpinPage(this.processNumber, vpn);
 				    return amountRead;
 				}
 				//we should only read until the end of the page, as we do 
				// not know where the next physical page may be. Will have
				//to recalculate the virtual address after that. 
				int amount = Math.min(length, pageSize - length);
				System.arraycopy(memory, paddr, data, offset, amount);
				amountRead = amountRead + amount;
				vaddr = vaddr + amount;
				length = length - amount;
				offset = offset + amount;
    	 
				for(int i = 0; i < Machine.processor().getTLBSize(); i++){
					TranslationEntry te = Machine.processor().readTLBEntry(i);
					if(te.vpn == vpn){
						te.used = true;
						Machine.processor().writeTLBEntry(i, te);
                 	}
				}
				
				VMKernel.getPageFromMemory(this.processNumber, vpn).used = true;
         }
         //unpin the page:
         VMKernel.unpinPage(this.processNumber, vpn);
 		}catch(IOException e){
			Lib.debug(dbgProcess, "IOException in getPhysicalAddress: " + e);
		}
 		Lib.debug(dbgVM, "----- exiting readVirtualMemory");
         return amountRead;
    }


    //Group 7
    //Method to handle a TLB miss
    public void handleTLBMiss(Processor processor){
    	
         int addr = Machine.processor().readRegister(Processor.regBadVAddr);
         int vpn = Processor.pageFromAddress(addr);
         Lib.debug(dbgVM, "handling TLBMiss for addr: " + addr + " vpn: " + vpn);
         //according to assignment, the key for the inverted page table is a combination of the
         //process id and virtual page number....so the format of the key will be "processID:pageNumber"
         String key = VMKernel.createKey(processNumber,vpn);
                 
             //page entry not in the page table
         	if (!VMKernel.isInMemory(processNumber, vpn)){
         		Lib.debug(dbgProcess, "page fault");
         		handlePageFault(vpn);
         	}
             Integer ppn = VMKernel.getPageFromMemory(key).ppn;
             addTLBEntry(vpn, ppn);
         
    }
    
    protected void addTLBEntry(int vpn, int ppn){
    	boolean hasEmptyEntry = false;
    	for(int i = 0; i < Machine.processor().getTLBSize(); i++){
            TranslationEntry te = Machine.processor().readTLBEntry(i);
            if(te == null || !te.valid){
                hasEmptyEntry = true;
                TranslationEntry newTe = new TranslationEntry(vpn,ppn,true,false,false,false);
                Machine.processor().writeTLBEntry(i, newTe);
                break;
            }
        }
        //if an empty TLB entry was not found then need to replace an existing one
        if(!hasEmptyEntry){
            Random random = new Random();
            int replacedEntry = random.nextInt(Machine.processor().getTLBSize());
             TranslationEntry newTe = new TranslationEntry(vpn,ppn,true,false,false,false);
            Machine.processor().writeTLBEntry(replacedEntry,newTe);
        }
    }
    
    //Group 7
    // function to handle the page fault. 
    private void handlePageFault(int vpn){
    	Lib.debug(dbgProcess, "---In handlePageFault ");
    	pageFaults++;
    	//first find if there is room in memory to swap the page in:
    	if(VMKernel.getFreeMemorySize() == 0){  //no room, swap out
    		try{
    			VMKernel.swapPageOut();
    		}catch(IOException e){
    			this.handleSyscall(1, 1, 0, 0, 0);
    		}
    	}
    	
    	//now swap page in if it is in swap. 
    	if(VMKernel.isPageInSwap(this.processNumber, vpn)){
    		Lib.debug(dbgProcess, "swap");
    		try{
    			VMKernel.swapPageIn(this.processNumber, vpn);
    		}catch(IOException e){
    			this.handleSyscall(1, 1, 0, 0, 0);
    			
    		}
    	}else{
    		/// need to load the page from disk
    		Lib.debug(dbgProcess, "loading page from the disk, not in swap");
    		int ppn = (Integer)VMKernel.getFreePage();
    		if(ppn < 0){
    			this.handleSyscall(1, 1, 0, 0, 0); //kill process if we cannot get memory
    		}
    		//Should have gotten a TE at this point
            Lib.debug(dbgProcess, "Adding page to memory with pid: " + this.processNumber
            			+ " and vpn: " + vpn + " and ppn: " + ppn);
            
            //Make this check now before assigning to the invPageTable
            boolean isSection = pageSectionMap.containsKey(vpn);
            if(isSection == false){
    			//Ensure that we don't over-allocate stack pages
    			totalStackPages++;
    			Lib.assertTrue(totalStackPages < stackPages + 1);
            }
            
            String key = VMKernel.createKey(this.processNumber, vpn);
    		TranslationEntry newTe = new TranslationEntry(vpn,ppn, true,false,false,false);
            VMKernel.addPageToMemory(key, newTe);
            
    		if(isSection){
    			//it is an executable, load it in
    			int coffSecNb = pageSectionMap.get(vpn);
        		CoffSection section = coff.getSection(coffSecNb);
        		VMKernel.getPageFromMemory(key).readOnly = true;
        		 //load section based on vpn, if initialized
                if(section.isInitialzed()){
                	//get the page number within this segment
                	Lib.debug(dbgVM, "Loading executable");
                	int pageNumber = vpn - section.getFirstVPN();
                	VMKernel.pinPage(this.processNumber, vpn);
                	section.loadPage(pageNumber, ppn);
                	VMKernel.unpinPage(this.processNumber, vpn);
                }
    		}
      	}
    	Lib.debug(dbgVM, "---Exiting handlePageFault");
    }
        
    protected int getPhysicalAddress(int vaddr)throws IOException{
    	//first get the virtual page which contains this address:
    	int vpn = Processor.pageFromAddress(vaddr); 
    	//now get the physical page:
    	int ppn = VMKernel.getPageFromMemory(this.processNumber, vpn).ppn;
		//now get the offset in the page:
		int offset = vaddr % pageSize;
		//now get the processor to calculate the physical address and return it:
		return Processor.makeAddress(ppn, offset);
    }
    
    public static void selfTest(){
    	VMProcess up = new VMProcess();
    	int vaddress = 50;
    	int ppn;
    	
    	Lib.debug(dbgVM, "Created process");
    	//process.handlePageFault(0);
    	byte[] dataOut = Lib.bytesFromInt(23);
    	Lib.debug(dbgProcess, "bytes to write: " + dataOut.length + " data to write: 23" );
    	int bytesProcessed = up.writeVirtualMemory(vaddress, dataOut); 
    	Lib.debug(dbgProcess, "wrote this many bytes: " + bytesProcessed);
    	//check if we can read it now
    	
    	byte[] dataIn = new byte[dataOut.length];
    	int bytesRead = up.readVirtualMemory(vaddress, dataIn);
    	Lib.debug(dbgProcess, "read this many bytes: " + bytesRead);
    	int result = Lib.bytesToInt(dataIn, 0);
    	Lib.debug(dbgProcess, "data read in: " + result);
    	
    	int freePages = VMKernel.getFreeMemorySize();
    	for(int i = 0; i < freePages; i++){
    		TranslationEntry te = new TranslationEntry(i, VMKernel.getFreePage(), false, false, false, false);
    		VMKernel.addPageToMemory("2:" + i, te);
    	}
    	
    	for(int i=3; i<5; i++){
	    	vaddress = Processor.makeAddress(i, 3);
	    	Lib.debug(dbgVM, "the virtual address is: " + vaddress);
	    	
	    	//now the memory should be full, try writing to it again. 
	    	dataOut = Lib.bytesFromInt(i);
	    	Lib.debug(dbgProcess, "bytes to write: " + dataOut.length + " data to write: " + i);
	    	bytesProcessed = up.writeVirtualMemory(vaddress, dataOut); 
	    	Lib.debug(dbgProcess, "wrote this many bytes: " + bytesProcessed);
	    	//check if we can read it now
	    	
	    	dataIn = new byte[dataOut.length];
	    	bytesRead = up.readVirtualMemory(vaddress, dataIn);
	    	Lib.debug(dbgProcess, "read this many bytes: " + bytesRead);
	    	result = Lib.bytesToInt(dataIn, 0);
	    	Lib.debug(dbgProcess, "data read in: " + result);
    	}
    	
    	
    	
    	Machine.terminate();
    }
    //group 7
    private static int totalProcesses;
    private int totalStackPages;
    private static Semaphore mutex = new Semaphore(1);
    private HashMap<Integer,Integer> pageSectionMap = new HashMap<Integer,Integer>();


    protected static final int pageSize = Processor.pageSize;
    protected static final char dbgProcess = 'a';
    protected static final char dbgVM = 'v';
    protected int pageFaults = 0;
    
}
